package simpledb.execution;

import simpledb.common.Type;
import simpledb.storage.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Knows how to compute some aggregate over a set of IntFields.
 */
public class IntegerAggregator implements Aggregator {

    private static final long serialVersionUID = 1L;
    private int gbfield;
    private Type gbfieldtype;
    private int afield;
    private Op what;
    private Map<Field, int[]> groups;
    private int[] nogroup;

    /**
     * Aggregate constructor
     *
     * @param gbfield     the 0-based index of the group-by field in the tuple, or
     *                    NO_GROUPING if there is no grouping
     * @param gbfieldtype the type of the group by field (e.g., Type.INT_TYPE), or null
     *                    if there is no grouping
     * @param afield      the 0-based index of the aggregate field in the tuple
     * @param what        the aggregation operator
     */

    public IntegerAggregator(int gbfield, Type gbfieldtype, int afield, Op what) {
        // some code goes here
        this.gbfield = gbfield;
        this.gbfieldtype = gbfieldtype;
        this.afield = afield;
        this.what = what;
        if (gbfieldtype != null) {
            this.groups = new HashMap<>();
        } else {
            this.nogroup = new int[4];
            nogroup[2] = Integer.MAX_VALUE;
            nogroup[3] = Integer.MIN_VALUE;
        }
    }

    /**
     * Merge a new tuple into the aggregate, grouping as indicated in the
     * constructor
     *
     * @param tup the Tuple containing an aggregate field and a group-by field
     */
    public void mergeTupleIntoGroup(Tuple tup) {
        IntField intField = (IntField) tup.getField(afield);
        if (gbfieldtype == null) {
            nogroup[0]++;
            nogroup[1] += intField.getValue();
            nogroup[2] = Math.min(nogroup[2], intField.getValue());
            nogroup[3] = Math.max(nogroup[3], intField.getValue());
        } else {
            Field identifier = tup.getField(gbfield);
            this.groups.putIfAbsent(identifier, new int[]{0, 0, Integer.MAX_VALUE, Integer.MIN_VALUE});

            int[] tmp = this.groups.get(identifier);
            tmp[0]++;
            tmp[1] += intField.getValue();
            tmp[2] = Math.min(tmp[2], intField.getValue());
            tmp[3] = Math.max(tmp[3], intField.getValue());

        }
    }

    /**
     * Create a OpIterator over group aggregate results.
     *
     * @return a OpIterator whose tuples are the pair (groupVal, aggregateVal)
     * if using group, or a single (aggregateVal) if no grouping. The
     * aggregateVal is determined by the type of aggregate specified in
     * the constructor.
     */
    public OpIterator iterator() {
        // some code goes here
        List<Tuple> tuples = new ArrayList<>();
        TupleDesc td;
        if (this.gbfieldtype != null) {
            td = new TupleDesc(new Type[]{this.gbfieldtype, Type.INT_TYPE});
            for (Map.Entry<Field, int[]> entry : this.groups.entrySet()) {
                Tuple a = new Tuple(td);
                a.setField(0, entry.getKey());
                int[] v = entry.getValue();
                switch (this.what) {
                    case MIN:
                        a.setField(1, new IntField(v[2]));
                        break;
                    case MAX:
                        a.setField(1, new IntField(v[3]));
                        break;
                    case SUM:
                        a.setField(1, new IntField(v[1]));
                        break;
                    case AVG:
                        a.setField(1, new IntField(v[1] / v[0]));
                        break;
                    case COUNT:
                        a.setField(1, new IntField(v[0]));
                        break;
                }
                tuples.add(a);
            }

        } else {
            td = new TupleDesc(new Type[]{Type.INT_TYPE});
            Tuple a = new Tuple(td);
            switch (this.what) {
                case MIN:
                    a.setField(0, new IntField(nogroup[2]));
                    break;
                case MAX:
                    a.setField(0, new IntField(nogroup[3]));
                    break;
                case SUM:
                    a.setField(0, new IntField(nogroup[1]));
                    break;
                case AVG:
                    a.setField(0, new IntField(nogroup[1] / nogroup[0]));
                    break;
                case COUNT:
                    a.setField(0, new IntField(nogroup[0]));
                    break;
            }
            tuples.add(a);
        }
        return new TupleIterator(td, tuples);
    }

}
